from turtle import forward
import torch
from torch import nn
from ..module.activation import act_layers


class BaseClassifier(nn.Module):
    """build base classifier without activation,
    we don't use nn.Linner in this classfiler

    Args:
        nn (_type_): _description_
    """

    def __init__(self, last_channel,nclasses,activation='Softmax', use_linear=False) -> None:
        super().__init__()
        self.nclasses = nclasses
        # self.in_channels = in_channels
        # self.classifier = nn.Conv2d(
        #     in_channels, out_channels, kernel_size=1, stride=1)
        self.use_linear = use_linear
        self.activation = act_layers(activation)
        if use_linear is False:
            self.linner = nn.Sequential(
                nn.AdaptiveAvgPool2d((1,1)),
                nn.Conv2d(last_channel,nclasses,kernel_size=1)
            )
        else:
            self.linner = nn.Sequential(
                nn.AdaptiveAvgPool2d((1,1)),
                nn.Flatten(),
                nn.Linear(last_channel,4096),
                nn.Linear(4096,nclasses)
            )
    def forward(self, x):
        x = self.linner(x)
        x = x.view(x.size(0),-1)
        x = self.activation(x)

        return x

class GlobalConvLinear(nn.Module):
    def __init__(self,in_channels,out_channels,kernel_size,stride=1) -> None:
        super().__init__()
        self.conv = nn.Conv2d(in_channels,in_channels,kernel_size,stride=stride)
        self.linner = nn.Conv2d(in_channels,out_channels,kernel_size=1,stride=1)

    def forward(self,x):
        x = self.conv(x)
        x = self.linner(x)
        return x
